Skip to content

feat(trainer): enable torch.compile by default#2489

Open
samsja wants to merge 1 commit into
mainfrom
feat/enable-torch-compile-by-default
Open

feat(trainer): enable torch.compile by default#2489
samsja wants to merge 1 commit into
mainfrom
feat/enable-torch-compile-by-default

Conversation

@samsja
Copy link
Copy Markdown
Member

@samsja samsja commented May 13, 2026

Summary

  • Default ModelConfig.compile to CompileConfig() instead of None, so apply_compile runs by default and each transformer block is wrapped with torch.compile.
  • fullgraph still defaults to False. Configs that explicitly set [model.compile] (e.g. configs/nemotron_4node/rl.toml, configs/hendrycks_math/sanity.toml) continue to override as before.

🤖 Generated with Claude Code


Note

Medium Risk
Changes default training behavior by enabling torch.compile unless explicitly disabled, which can impact performance, debugging, and runtime compatibility for some models/environments.

Overview
Enables model compilation by default by changing ModelConfig.compile’s default from None to CompileConfig(), so configurations that don’t specify model.compile will now run the model through torch.compile automatically.

Existing configs that explicitly set model.compile continue to override this default.

Reviewed by Cursor Bugbot for commit 0f2fc1f. Bugbot is set up for automated code reviews on this repo. Configure here.

S1ro1
S1ro1 previously approved these changes May 13, 2026
Copy link
Copy Markdown
Collaborator

@S1ro1 S1ro1 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

big+++

@samsja samsja marked this pull request as ready for review May 13, 2026 23:54
Copy link
Copy Markdown

@cursor cursor Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit f95f67a. Configure here.

Comment thread packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py Outdated
Default `ModelConfig.compile` to `CompileConfig()` so transformer
blocks are compiled by default. Configs that explicitly set
`[model.compile]` continue to override as before.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@samsja samsja force-pushed the feat/enable-torch-compile-by-default branch from f95f67a to 0f2fc1f Compare May 14, 2026 04:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants